mod statistics;
mod statistics_polars;
mod statm;

use anyhow::Result;
use clap::Parser;
use polars::frame::DataFrame;
use polars::io::SerWriter as _;
use polars::prelude::{CsvWriter, ParquetWriter};
use polars_excel_writer::PolarsXlsxWriter;
use statistics::Statistics;
use statm::MemInfo;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::{thread::sleep, time::Duration};
use strum::VariantNames;
use strum_macros::{EnumString, VariantNames};
use tracing::{error, info, warn};
extern crate chrono;
use chrono::{Datelike, Local, Timelike};

#[derive(Debug, Clone, VariantNames, EnumString, strum_macros::Display)]
enum SaveFormat {
    CSV,
    Excel,
    Parquet,
}

macro_rules! clap_enum_variants {
    ($e: ty) => {{
        use clap::builder::TypedValueParser;
        clap::builder::PossibleValuesParser::new(<$e>::VARIANTS).map(|s| s.parse::<$e>().unwrap())
    }};
}

#[derive(Parser, Debug)]
#[command( version, about, long_about = None)]
struct Args {
    #[arg(long)]
    pid: u64,
    #[arg(long, default_value_t = 4096)]
    pagesize: u64,
    #[arg(long, default_value_t = 1)]
    duration: u8,
    #[arg(long, short, value_parser = clap_enum_variants!(SaveFormat))]
    save_fmt: Option<SaveFormat>,
}

fn main() {
    tracing_subscriber::fmt::init();

    let args = Args::parse();
    info!("Looking at pid: {}", args.pid);
    let mut statistics = Statistics::new();

    let running = Arc::new(AtomicBool::new(true));
    let r = running.clone();

    ctrlc::set_handler(move || {
        info!("Ctrl-C received, exiting");
        r.store(false, Ordering::SeqCst);
    })
    .expect("Error setting Ctrl-C handler");

    loop {
        if !running.load(Ordering::SeqCst) {
            if let Some(save_fmt) = args.save_fmt {
                info!("Saving statistics to file in {save_fmt} format");
                match statistics.to_dataframe() {
                    Ok(mut df) => match save_file(&mut df, save_fmt) {
                        Ok(path) => {
                            info!("Statistics saved successfully: {}", path);
                        }
                        Err(e) => {
                            error!("Error saving statistics to file: {e}");
                            panic!();
                        }
                    },
                    Err(e) => {
                        error!("Error saving statistics to file: {e}");
                        panic!();
                    }
                }
            }
            break;
        }
        match MemInfo::new(args.pid, args.pagesize) {
            Ok(meminfo) => {
                statistics.add(&meminfo);
                statistics.display();
            }
            Err(e) => {
                warn!("Error getting meminfo: {e}");
                if statistics.is_empty() {
                    error!("Process with pid {} not exist", args.pid);
                    panic!();
                } else {
                    warn!("Process with pid {} terminated", args.pid);
                    running.store(false, Ordering::SeqCst);
                }
            }
        }
        sleep(Duration::from_secs(args.duration.into()));
    }
}

fn save_file(df: &mut DataFrame, save_fmt: SaveFormat) -> Result<String> {
    let local = Local::now();
    let filename = format!(
        "meminfo_{}_{}_{}_{}_{}_{}",
        local.year(),
        local.month(),
        local.day(),
        local.hour(),
        local.minute(),
        local.second()
    );

    match save_fmt {
        SaveFormat::CSV => {
            let path = format!("{}.csv", filename);
            let mut file = std::fs::File::create(&path)?;
            CsvWriter::new(&mut file).finish(df)?;
            Ok(path)
        }
        SaveFormat::Excel => {
            let path = format!("{}.xlsx", filename);
            let mut xlsx_writer = PolarsXlsxWriter::new();
            xlsx_writer.write_dataframe(df)?;
            xlsx_writer.save(&path)?;
            Ok(path)
        }
        SaveFormat::Parquet => {
            let path = format!("{}.parquet", filename);
            let mut file = std::fs::File::create(&path)?;
            ParquetWriter::new(&mut file).finish(df)?;
            Ok(path)
        }
    }
}
